import openslide
import os
import glob
import time
from PIL import Image  # 导入 PIL.Image，尽管 openslide 会返回它

# --- 1. 配置路径 ---

# 指定包含 SVS 文件的目录 (与您之前的脚本一致)
# wsi_directory = "/mnt/data1/zzy/LNM/CRCpro-wsi/2025_10_ALL/"
wsi_directory = "/mnt/data1/zzy/LNM/slices"

# !!! 指定您希望保存缩略图的新目录 !!!
thumbnail_output_directory = (
    "/mnt/data1/zzy/LNM/slices_thumbnails/"
)

# 缩略图的文件格式 (png 或 jpg)
thumbnail_format = "png"

# --- 2. 确保输出目录存在 ---
print(f"缩略图将保存到: {thumbnail_output_directory}")
os.makedirs(thumbnail_output_directory, exist_ok=True)

# --- 3. 查找所有 SVS 文件 ---
svs_files = glob.glob(os.path.join(wsi_directory, "*.svs"))

if not svs_files:
    print(f"在目录 {wsi_directory} 中未找到任何 .svs 文件。")
    exit()

print(f"找到了 {len(svs_files)} 个 .svs 文件，开始提取缩略图...")
print("=============================================")

start_time = time.time()
success_count = 0
fail_count = 0

# --- 4. 遍历和处理每个文件 ---
for i, svs_file_path in enumerate(svs_files):
    file_name = os.path.basename(svs_file_path)
    # 构建输出文件名，例如 "image1.svs" -> "image1_thumbnail.png"
    base_name = os.path.splitext(file_name)[0]
    output_filename = f"{base_name}_thumbnail.{thumbnail_format}"
    output_path = os.path.join(thumbnail_output_directory, output_filename)

    print(f"[{i+1}/{len(svs_files)}] 正在处理: {file_name}")

    try:
        slide = openslide.OpenSlide(svs_file_path)

        # 检查 'thumbnail' 是否存在于关联图像中
        if "thumbnail" in slide.associated_images:
            # 提取缩略图 (这是一个 PIL.Image 对象)
            thumbnail_image = slide.associated_images["thumbnail"]

            # 保存缩略图
            # 如果缩略图是 RGBA (带透明度)，转换为 RGB 以便保存为 JPG
            if thumbnail_format.lower() == "jpg" and thumbnail_image.mode == "RGBA":
                thumbnail_image = thumbnail_image.convert("RGB")

            thumbnail_image.save(output_path)
            print(f"  > 成功保存缩略图到: {output_filename}")
            success_count += 1

        else:
            print("  > 未找到 'thumbnail' 关联图像。")
            fail_count += 1

        slide.close()

    except Exception as e:
        print(f"  ! 处理 {file_name} 时发生错误: {e}")
        fail_count += 1
        if "slide" in locals() and slide:
            slide.close()

# --- 5. 结束 ---
end_time = time.time()
print("\n=============================================")
print(f"所有文件处理完毕。")
print(f"  - 成功提取: {success_count} 个")
print(f"  - 失败/跳过: {fail_count} 个")
print(f"  - 总耗时: {end_time - start_time:.2f} 秒。")
print(f"缩略图已保存到: {thumbnail_output_directory}")
